import cv2
import os
from shutil import rmtree


def empty_dir(path):
    """ empty specified dir """
    if os.path.exists(path):
        rmtree(path)
    os.mkdir(path)


if __name__ == '__main__':
    dirtype = ("train", "label", "test")
    strpwd = os.getcwd()
    print("pwd = %s" % strpwd)


    loadDir = r"E:\xianhuang0818\xhsi_seg20210113\img_jin"
    loadDir_mask = r"E:\xianhuang0818\xhsi_seg20210113\masks"
    saveDir_1 = r"E:\xianhuang0818\xhsi_seg20210113\img_jin_equ"
    saveDir_2 = r"E:\xianhuang0818\xhsi_seg20210113\img_jin_mask"
    # loadDir = r"E:\xianhuang0818\xhsi_seg20210113\img_yin"
    # loadDir_mask = r"E:\xianhuang0818\xhsi_seg20210113\masks"
    # saveDir_1 = r"E:\xianhuang0818\xhsi_seg20210113\img_yin_equ"
    # saveDir_2 = r"E:\xianhuang0818\xhsi_seg20210113\img_yin_mask"

    # loadDir = r"E:\xianhuang0818\xhsi_seg20201226\img_jin"
    # loadDir_mask = r"E:\xianhuang0818\xhsi_seg20201226\masks"
    # saveDir_1 = r"E:\xianhuang0818\xhsi_seg20201226\img_jin_equ"
    # saveDir_2 = r"E:\xianhuang0818\xhsi_seg20201226\img_jin_mask"


    empty_dir(saveDir_1)
    empty_dir(saveDir_2)
    imgNames = os.listdir(loadDir)
    # print(imgNames)
    for j in range(0, len(imgNames)):
        title = imgNames[j].split(".")[0]
        print(imgNames[j])
        imgPath = os.path.join(loadDir, imgNames[j])
        img = cv2.imread(imgPath, 1)
        cv2.imshow("bgr", cv2.resize(img, (600,600)))
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        equ = cv2.equalizeHist(gray)
        writePath = os.path.join(saveDir_1, '%s.png'%title)
        cv2.imwrite(writePath, equ)
        cv2.imshow("equ", cv2.resize(equ, (600,600)))

        imgPath = os.path.join(loadDir_mask, imgNames[j])
        img = cv2.imread(imgPath, 1)
        writePath = os.path.join(saveDir_2, '%s.png'%title)
        cv2.imwrite(writePath, img)
        cv2.imshow("mask", cv2.resize(img, (600,600)))
        cv2.waitKey(100)



